import os
import torch
import tifffile
import numpy as np
import json
from tqdm import tqdm

from embedtrack.infer.inference import init_model
from embedtrack.utils.clustering import Cluster

# ========== 用户需填写以下路径 ==========
MODEL_PATH = 'your_model.pth'  # 模型权重路径
CONFIG_PATH = 'your_config.json'  # 配置文件路径
INPUT_FOLDER = 'your_input_folder'  # 输入图片文件夹
OUTPUT_FOLDER = 'your_output_folder'  # 输出mask文件夹

# ========== 加载配置 ==========
with open(CONFIG_PATH, 'r') as f:
    config = json.load(f)

model_class = config['model_dict']['name']
input_channels = config['model_dict']['kwargs']['input_channels']
n_classes = config['model_dict']['kwargs']['n_classes']
grid_y = config['grid_dict']['grid_y']
grid_x = config['grid_dict']['grid_x']
pixel_y = config['grid_dict']['pixel_y']
pixel_x = config['grid_dict']['pixel_x']
crop_size = config['train_dict']['crop_size']
min_mask_size = config['train_dict'].get('min_mask_size', 10)

# ========== 初始化模型 ==========
model_dict = {
    'kwargs': {
        'input_channels': input_channels,
        'n_classes': n_classes,
    }
}
model_configs = {
    'model_class': model_class,
    'model_cktp_path': MODEL_PATH
}
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = init_model(model_dict, model_configs)
model = model.to(device)
model.eval()

# ========== 初始化Cluster ==========
cluster = Cluster(grid_y, grid_x, pixel_y, pixel_x)
cluster = cluster.to(device)

# ========== 创建输出文件夹 ==========
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# ========== 推理每张图片 ==========
img_list = [f for f in os.listdir(INPUT_FOLDER) if f.lower().endswith(('.tif', '.tiff', '.png', '.jpg', '.jpeg'))]

for img_name in tqdm(img_list, desc='推理中'):
    img_path = os.path.join(INPUT_FOLDER, img_name)
    image = tifffile.imread(img_path)
    if image.ndim == 2:
        image = image[np.newaxis, ...]  # (1, H, W)
    elif image.ndim == 3 and image.shape[0] != input_channels:
        image = np.transpose(image, (2, 0, 1))  # (C, H, W)
    input_tensor = torch.from_numpy(image).float().unsqueeze(0).to(device)  # (1, C, H, W)
    with torch.no_grad():
        pred = model(input_tensor)
        if isinstance(pred, (tuple, list)):
            pred = pred[0]  # 取分割输出
        mask = cluster.cluster_pixels(pred.squeeze(0), n_sigma=2, min_obj_size=min_mask_size)
    mask = mask.cpu().numpy().astype('uint16')
    tifffile.imwrite(os.path.join(OUTPUT_FOLDER, os.path.splitext(img_name)[0] + '_mask.tif'), mask)

print('所有图片分割完成，结果已保存到:', OUTPUT_FOLDER)
